"""
NS-SEM Solver Extension Implementation

Main extension class that integrates the Navier-Stokes Spectral Element Method
solver with PETSc-Kokkos framework into GSM.
"""

from typing import Dict, List, Any
from pathlib import Path
import os
import logging

from gcr_solver_manager.extensions.base_extension import BaseExtension
from gcr_solver_manager.extensions.base_builder import BaseBuilder
from gcr_solver_manager.extensions.base_runner import BaseRunner
from gcr_solver_manager.extensions.base_validator import BaseValidator

from .builder import NSsemBuilder
from .runner import NSsemRunner
from .validator import NSsemValidator

logger = logging.getLogger(__name__)


class NSsemSolverExtension(BaseExtension):
    """
    Extension for NS-SEM (Navier-Stokes Spectral Element Method) solver.
    
    Provides support for building and running NS-SEM solvers with PETSc-Kokkos
    integration for high-performance fluid dynamics simulations.
    """
    
    def __init__(self):
        super().__init__("ns-sem-solver", "1.0.0")
        self._builder = None
        self._runner = None
        self._validator = None
    
    @property
    def description(self) -> str:
        """Get extension description."""
        return ("Navier-Stokes Spectral Element Method solver with PETSc-Kokkos "
                "integration providing high-performance fluid dynamics simulations "
                "with tensor matrix optimization and multi-cavity support")
    
    @property
    def supported_frameworks(self) -> List[str]:
        """Get supported framework names."""
        return [
            "ns-sem-solver",
            "ns-sem-petsc",
            "spectral-element-method",
            "navier-stokes-sem",
            "ns-sem-gcr",
            "ns-sem-gmres"
        ]
    
    @property
    def config_schema(self) -> Dict[str, Any]:
        """Get JSON schema for configuration validation."""
        return {
            "$schema": "http://json-schema.org/draft-07/schema#",
            "type": "object",
            "title": "NS-SEM Solver Configuration Schema",
            "properties": {
                "Problem": {
                    "type": "object",
                    "properties": {
                        "problem_type": {
                            "type": "string",
                            "enum": ["navier_stokes", "helmholtz"],
                            "default": "navier_stokes"
                        },
                        "spectral_order": {
                            "type": "integer",
                            "minimum": 2,
                            "maximum": 20,
                            "default": 8
                        },
                        "grid_nx": {"type": "integer", "minimum": 1, "default": 64},
                        "grid_ny": {"type": "integer", "minimum": 1, "default": 64},
                        "grid_nz": {"type": "integer", "minimum": 1, "default": 32},
                        "reynolds_number": {
                            "type": "number",
                            "minimum": 1.0,
                            "default": 1000.0
                        },
                        "time_step": {
                            "type": "number",
                            "minimum": 1e-6,
                            "default": 0.001
                        },
                        "final_time": {
                            "type": "number",
                            "minimum": 0.0,
                            "default": 1.0
                        }
                    },
                    "required": ["problem_type", "spectral_order"]
                },
                "Solver": {
                    "type": "object",
                    "properties": {
                        "type": {
                            "type": "string",
                            "enum": ["gcr", "ca-gcr", "gmres", "bca-gmres", "cg", "pcg"],
                            "default": "gcr"
                        },
                        "max_iterations": {
                            "type": "integer",
                            "minimum": 1,
                            "maximum": 10000,
                            "default": 1000
                        },
                        "convergence_tol": {
                            "type": "number",
                            "minimum": 1e-15,
                            "maximum": 1e-3,
                            "default": 1e-8
                        },
                        "restart_length": {
                            "type": "integer",
                            "minimum": 1,
                            "maximum": 1000,
                            "default": 50
                        },
                        "use_petsc": {
                            "type": "boolean",
                            "default": true
                        },
                        "tensor_format": {
                            "type": "string",
                            "enum": ["19_diagonal", "full_tensor"],
                            "default": "19_diagonal"
                        }
                    },
                    "required": ["type"]
                },
                "System": {
                    "type": "object",
                    "properties": {
                        "gpu_backend": {
                            "type": "string",
                            "enum": ["cuda", "hip", "none"],
                            "default": "cuda"
                        },
                        "communication": {
                            "type": "string",
                            "enum": ["nccl", "rccl", "mpi", "none"],
                            "default": "nccl"
                        },
                        "solver_framework": {
                            "type": "string",
                            "enum": ["petsc", "native"],
                            "default": "petsc"
                        }
                    }
                },
                "MPI": {
                    "type": "object",
                    "properties": {
                        "total_processes": {
                            "type": "integer",
                            "minimum": 1,
                            "default": 4
                        },
                        "nprocx": {"type": "integer", "minimum": 1, "default": 2},
                        "nprocy": {"type": "integer", "minimum": 1, "default": 2},
                        "nprocz": {"type": "integer", "minimum": 1, "default": 1}
                    }
                },
                "TensorMatrix": {
                    "type": "object",
                    "properties": {
                        "matrix_type": {
                            "type": "string",
                            "enum": ["19_diagonal", "27_point"],
                            "default": "19_diagonal"
                        },
                        "assembly_method": {
                            "type": "string",
                            "enum": ["local", "global"],
                            "default": "local"
                        },
                        "compression": {
                            "type": "boolean",
                            "default": true
                        }
                    }
                },
                "MultiCavity": {
                    "type": "object",
                    "properties": {
                        "enable_multi_cavity": {
                            "type": "boolean",
                            "default": false
                        },
                        "cavity_spacing": {
                            "type": "number",
                            "minimum": 0.01,
                            "default": 0.1
                        },
                        "boundary_coupling": {
                            "type": "string",
                            "enum": ["weak", "strong"],
                            "default": "weak"
                        }
                    }
                }
            },
            "required": ["Problem", "Solver"]
        }
    
    def can_handle_config(self, config_path: str) -> bool:
        """
        Check if this extension can handle the given configuration.
        
        Args:
            config_path: Path to configuration file
            
        Returns:
            True if extension can handle this configuration
        """
        config_path_obj = Path(config_path)
        
        # Check path patterns
        path_str = str(config_path_obj).lower()
        
        # High confidence matches
        if "ns-sem" in path_str or "ns_sem" in path_str:
            logger.debug(f"NS-SEM path pattern matched: {config_path}")
            return True
        
        if "navier-stokes" in path_str or "navier_stokes" in path_str:
            logger.debug(f"Navier-Stokes pattern matched: {config_path}")
            return True
        
        if "spectral-element" in path_str or "spectral_element" in path_str:
            logger.debug(f"Spectral element pattern matched: {config_path}")
            return True
        
        # Medium confidence matches
        if config_path_obj.name.startswith("ns_sem_"):
            return True
        
        # Check if config file exists and contains NS-SEM markers
        if config_path_obj.exists():
            try:
                content = config_path_obj.read_text().lower()
                ns_sem_markers = [
                    "ns-sem",
                    "ns_sem",
                    "navier_stokes",
                    "spectral_order",
                    "reynolds_number",
                    "tensor_format",
                    "19_diagonal",
                    "multi_cavity",
                    "tensor_mat",
                    "spectral element"
                ]
                
                marker_count = sum(1 for marker in ns_sem_markers if marker in content)
                if marker_count >= 3:  # Need at least 3 markers for confidence
                    logger.debug(f"Content analysis matched: {config_path} (markers: {marker_count})")
                    return True
                    
            except Exception as e:
                logger.debug(f"Error reading config file {config_path}: {e}")
        
        return False
    
    def get_builder(self) -> BaseBuilder:
        """Get the builder instance for this extension."""
        if self._builder is None:
            self._builder = NSsemBuilder("ns-sem-solver")
        return self._builder
    
    def get_runner(self) -> BaseRunner:
        """Get the runner instance for this extension."""
        if self._runner is None:
            self._runner = NSsemRunner("ns-sem-solver")
        return self._runner
    
    def get_validator(self) -> BaseValidator:
        """Get the validator instance for this extension.""" 
        if self._validator is None:
            self._validator = NSsemValidator("ns-sem-solver")
        return self._validator
    
    def initialize(self) -> None:
        """Initialize the extension."""
        logger.info("Initializing NS-SEM Solver extension")
        
        # Check environment
        validation_errors = self.validate_environment()
        if validation_errors:
            logger.warning(f"Environment validation warnings: {validation_errors}")
        
        # Initialize components
        try:
            self.get_builder()
            self.get_runner()
            self.get_validator()
            logger.info("NS-SEM Solver extension initialized successfully")
        except Exception as e:
            logger.error(f"Failed to initialize extension components: {e}")
            raise
    
    def cleanup(self) -> None:
        """Clean up extension resources."""
        logger.info("Cleaning up NS-SEM Solver extension")
        
        # Clean up component instances
        self._builder = None
        self._runner = None
        self._validator = None
    
    def validate_environment(self) -> List[str]:
        """
        Validate that the environment meets extension requirements.
        
        Returns:
            List of error messages (empty if valid)
        """
        errors = []
        
        # Check for PETSc
        petsc_dir = os.getenv('PETSC_DIR')
        if not petsc_dir:
            errors.append("PETSC_DIR environment variable not set")
        elif not Path(petsc_dir).exists():
            errors.append(f"PETSC_DIR path does not exist: {petsc_dir}")
        
        # Check for Kokkos
        kokkos_root = os.getenv('Kokkos_ROOT')
        if not kokkos_root:
            # Also check common installation paths
            common_paths = [
                Path("/usr/local/include/Kokkos_Core.hpp"),
                Path("/opt/kokkos/include/Kokkos_Core.hpp"),
            ]
            if not any(p.exists() for p in common_paths):
                errors.append("Kokkos not found (set Kokkos_ROOT or install to standard location)")
        elif not Path(kokkos_root).exists():
            errors.append(f"Kokkos_ROOT path does not exist: {kokkos_root}")
        
        # Check for MPI
        mpi_found = any(
            Path(f"/usr/bin/{cmd}").exists() or 
            Path(f"/usr/local/bin/{cmd}").exists() 
            for cmd in ["mpirun", "mpiexec"]
        )
        if not mpi_found:
            errors.append("MPI not found in standard locations")
        
        # Check for CMake
        cmake_found = any(
            Path(f"/usr/bin/cmake").exists(),
            Path(f"/usr/local/bin/cmake").exists()
        )
        if not cmake_found:
            errors.append("CMake not found in standard locations")
        
        return errors
    
    def get_commands(self) -> Dict[str, Any]:
        """Get custom CLI commands provided by this extension."""
        return {
            "ns-sem": {
                "help": "NS-SEM Solver framework commands",
                "subcommands": {
                    "build": {
                        "help": "Build NS-SEM solver framework",
                        "args": [
                            {"name": "--clean", "action": "store_true", "help": "Clean build"},
                            {"name": "--verbose", "action": "store_true", "help": "Verbose output"},
                            {"name": "--petsc", "action": "store_true", "help": "Enable PETSc integration"},
                            {"name": "--kokkos", "action": "store_true", "help": "Enable Kokkos acceleration"},
                        ]
                    },
                    "test": {
                        "help": "Run NS-SEM solver tests",
                        "args": [
                            {"name": "--tensor", "action": "store_true", "help": "Test tensor matrix operations"},
                            {"name": "--multi-cavity", "action": "store_true", "help": "Test multi-cavity functionality"},
                            {"name": "--petsc", "action": "store_true", "help": "Test PETSc integration"},
                            {"name": "--boundary", "action": "store_true", "help": "Test boundary communication"},
                        ]
                    },
                    "simulate": {
                        "help": "Run NS-SEM simulation",
                        "args": [
                            {"name": "--config", "type": str, "help": "Configuration file"},
                            {"name": "--reynolds", "type": float, "help": "Reynolds number"},
                            {"name": "--time-step", "type": float, "help": "Time step size"},
                            {"name": "--final-time", "type": float, "help": "Final simulation time"},
                        ]
                    },
                    "benchmark": {
                        "help": "Run performance benchmarks",
                        "args": [
                            {"name": "--compare", "action": "store_true", "help": "Compare with baseline"},
                            {"name": "--scaling", "action": "store_true", "help": "Scaling analysis"},
                        ]
                    }
                }
            }
        }
    
    def get_dependencies(self) -> List[str]:
        """Get list of other extensions this extension depends on."""
        return []  # No dependencies on other extensions